#include "fetcher.h"
#include "utils.h"
#include "cache_manager.h" 
#include <curl/curl.h>
#include <stdlib.h>
#include <string.h>
#include <stdbool.h> // <-- 增加包含

// libcurl 写回调函数
static size_t write_callback(void *contents, size_t size, size_t nmemb, void *userp) {
    size_t realsize = size * nmemb;
    fetch_task_arg_t *arg = (fetch_task_arg_t *)userp;
    
    memcpy(arg->buffer, contents, realsize);
    return realsize;
}

static void download_worker(void* arg) {
    fetch_task_arg_t *task_arg = (fetch_task_arg_t *)arg;
    CURL *curl;
    CURLcode res;
    bool success = false;

    char url[1024];
    // OCI digest 规范通常是 'sha256:' + hash, 但我们的 M1 使用的是 xxhash
    // 为了与 Registry API 兼容, 我们这里先硬编码为 sha256。
    // 生产级方案需要 M1 输出两种哈希, 或 Registry 支持 xxh64
    snprintf(url, sizeof(url), "%s/v2/%s/blobs/sha256:%s", 
             task_arg->fetcher->registry_url, 
             task_arg->fetcher->repository, 
             task_arg->hash); // 注意: 这里的 hash 需要是 SHA256 格式

    curl = curl_easy_init();
    if (curl) {
        task_arg->buffer = malloc(CHUNK_SIZE);

        curl_easy_setopt(curl, CURLOPT_URL, url);
        curl_easy_setopt(curl, CURLOPT_WRITEFUNCTION, write_callback);
        curl_easy_setopt(curl, CURLOPT_WRITEDATA, (void *)task_arg);
        curl_easy_setopt(curl, CURLOPT_FOLLOWLOCATION, 1L);

        res = curl_easy_perform(curl);
        if (res == CURLE_OK) {
            long http_code = 0;
            curl_easy_getinfo(curl, CURLINFO_RESPONSE_CODE, &http_code);
            if(http_code == 200) {
                success = true;
            } else {
                LOG_ERROR("下载失败, HTTP Code: %ld, URL: %s", http_code, url);
            }
        } else {
            LOG_ERROR("curl_easy_perform() failed: %s", curl_easy_strerror(res));
        }
        curl_easy_cleanup(curl);
    }
    
    pthread_mutex_lock(&task_arg->fetcher->inflight_lock);
    g_hash_table_remove(task_arg->fetcher->in_flight_hashes, task_arg->hash);
    pthread_mutex_unlock(&task_arg->fetcher->inflight_lock);

    task_arg->callback(task_arg, success);
}

fetcher_t* fetcher_create(const char* registry_url, const char* repository, int num_threads) {
    fetcher_t* f = (fetcher_t*)malloc(sizeof(fetcher_t));
    if (!f) return NULL;

    f->registry_url = strdup(registry_url);
    f->repository = strdup(repository);
    f->download_pool = thread_pool_create(num_threads, 128);
    f->in_flight_hashes = g_hash_table_new_full(g_str_hash, g_str_equal, g_free, NULL);
    pthread_mutex_init(&f->inflight_lock, NULL);
    
    curl_global_init(CURL_GLOBAL_DEFAULT);
    return f;
}

void fetcher_destroy(fetcher_t* f) {
    if (f) {
        thread_pool_destroy(f->download_pool);
        free(f->registry_url);
        free(f->repository);
        g_hash_table_destroy(f->in_flight_hashes);
        pthread_mutex_destroy(&f->inflight_lock);
        free(f);
        curl_global_cleanup();
    }
}

void fetcher_fetch_chunk(fetcher_t* f, const char* hash, void (*callback)(void*, bool), void* callback_arg) {
    pthread_mutex_lock(&f->inflight_lock);
    if (g_hash_table_contains(f->in_flight_hashes, hash)) {
        pthread_mutex_unlock(&f->inflight_lock);
        return;
    }
    // 注意：这里需要复制 hash 字符串，因为 g_hash_table_add 不会复制它
    g_hash_table_add(f->in_flight_hashes, g_strdup(hash));
    pthread_mutex_unlock(&f->inflight_lock);

    fetch_task_arg_t *task_arg = (fetch_task_arg_t*)malloc(sizeof(fetch_task_arg_t));
    task_arg->fetcher = f;
    task_arg->hash = strdup(hash);
    task_arg->callback = callback;
    task_arg->callback_arg = callback_arg;
    
    thread_pool_add(f->download_pool, download_worker, task_arg);
}